import os

import numpy as np
import pandas as pd
import seaborn as sns
from matplotlib import pyplot as plt
from matplotlib import rcParams
from sklearn.metrics.pairwise import cosine_similarity

from config import *


sns.set()

# To resolve the case of missing names
# fpaths = matplotlib.font_manager.findSystemFonts()
# fnames = [matplotlib.font_manager.get_font(f) for f in fpaths]
rcParams.update(
    {
        "font.family": "serif",
        "font.serif": "STSong, Times New Roman, SimSun, Times, Source Han Sans CN",
        "mathtext.fontset": "stix",
        "axes.unicode_minus": False,
    }
)


def get_statistics(extract_data, feature_name_list, is_washed=False):
    """Extract useful information from `extrac_data`
    TODO: remove is_washed
    """
    # Extract data statistics (7)
    extract_data_statistics = np.zeros((len(feature_name_list), 7))

    for i, feature in enumerate(feature_name_list):
        index = ~np.isnan(extract_data[i])
        missing_count = 0
        if not is_washed:
            missing_count = sum(
                [
                    feature_dict[feature]["total"] - np.sum(arr)
                    for arr in np.array_split(index, len(index) // 24)
                    if feature_dict[feature]["total"] > np.sum(arr)
                ]
            )
        total_count = np.sum(index)

        # 0. Total counts
        extract_data_statistics[i, 0] = total_count
        # 1. Missing value count
        extract_data_statistics[i, 1] = missing_count
        # 2. Mean value
        extract_data_statistics[i, 2] = np.mean(extract_data[i][index])
        # 3. Standard deviation
        extract_data_statistics[i, 3] = np.std(extract_data[i][index])
        # 4. Median
        extract_data_statistics[i, 4] = np.median(extract_data[i][index])
        # 5. Minimum value
        extract_data_statistics[i, 5] = np.min(extract_data[i][index])
        # 6. Maximum value
        extract_data_statistics[i, 6] = np.max(extract_data[i][index])

    return extract_data_statistics


def analyze_plot(extract_data, subdir=""):
    """Use `extract_data` to draw plots"""
    # Extract data statistics (7)
    hist_savedir = os.path.join(histogram_image_file_dir, subdir)
    box_savedir = os.path.join(boxplot_image_file_dir, subdir)
    os.makedirs(hist_savedir, exist_ok=True)
    os.makedirs(box_savedir, exist_ok=True)

    for i, feature in enumerate(feature_name_list):
        index = ~np.isnan(extract_data[i])

        logger.info(f"Rendering the histogram of feature <{feature}>")
        sns.displot(extract_data[i][index], kind="hist", stat="density", bins=100)
        plt.title(feature)
        plt.savefig(os.path.join(hist_savedir, f"{feature}_histogram.png"))
        plt.close()

        logger.info(f"Rendering the box-plot of feature  <{feature}>")
        sns.boxplot(data=extract_data[i][index], orient="v", color="skyblue")
        sns.stripplot(  # add scatter plot
            data=extract_data[i][index],
            color="orange",
            jitter=True,
            size=1.0,
        )
        plt.title(feature)
        plt.savefig(os.path.join(box_savedir, f"{feature}_boxplot.png"))
        plt.close()


def analyze_correlation(feature_data, feature_name_list, subdir=""):
    # Pearson correlation coefficient matrix
    pearson_cc_matrix = np.corrcoef(feature_data.T)
    correlate_savedir = os.path.join(correlation_image_file_dir, subdir)
    os.makedirs(correlate_savedir, exist_ok=True)

    logger.info(f"Rendering the heatmap of Pearson correlation coeffecient matrix")
    plt.figure(figsize=(12, 10))
    sns.heatmap(
        pearson_cc_matrix,
        xticklabels=feature_name_list,
        yticklabels=feature_name_list,
        linewidths=0.5,
    )
    plt.savefig(os.path.join(correlate_savedir, f"pearson_cca_heatmap.png"))
    plt.close()

    return pearson_cc_matrix


if __name__ == "__main__":
    feature_data = np.load(feature_data_file_path)
    extract_data = np.load(extract_data_file_path)

    washed_feature_data = np.load(washed_feature_data_file_path)
    washed_extract_data = np.load(washed_extract_data_file_path, allow_pickle=True)

    # Task 1
    analyze_plot(extract_data, subdir="original")
    extract_data_statistics = get_statistics(extract_data, feature_name_list)
    pd_extract_data_statistics = pd.DataFrame(
        extract_data_statistics, index=feature_name_list
    )
    analyze_plot(washed_extract_data, subdir="washed")
    washed_extract_data_statistics = get_statistics(
        washed_extract_data, washed_feature_name_list, True
    )
    washed_pd_extract_data_statistics = pd.DataFrame(
        washed_extract_data_statistics, index=washed_feature_name_list
    )

    # Task 2
    pearson_cc_matrix = analyze_correlation(
        feature_data, feature_name_list, subdir="original"
    )
    pd_pearson_cc_matrix = pd.DataFrame(pearson_cc_matrix, index=feature_name_list)
    washed_pearson_cc_matrix = analyze_correlation(
        washed_feature_data, washed_feature_name_list, subdir="washed"
    )
    washed_pd_pearson_cc_matrix = pd.DataFrame(
        washed_pearson_cc_matrix, index=washed_feature_name_list
    )

    # excel save
    with pd.ExcelWriter(excel_result_file_path, mode="w", engine="openpyxl") as writer:
        pd_extract_data_statistics.to_excel(
            writer,
            sheet_name="任务1",
            header=["总数", "缺失值个数", "均值", "标准差", "中位数", "最小值", "最大值"],
            index_label="指标量",
            float_format="%.5f",
        )
        pd_pearson_cc_matrix.to_excel(
            writer,
            sheet_name="任务2_皮尔逊相关系数",
            header=feature_name_list,
            index_label="指标量",
            float_format="%.5f",
        )
    with pd.ExcelWriter(
        excel_washed_result_file_path, mode="w", engine="openpyxl"
    ) as writer:
        washed_pd_extract_data_statistics.to_excel(
            writer,
            sheet_name="任务1",
            header=["总数", "缺失值个数", "均值", "标准差", "中位数", "最小值", "最大值"],
            index_label="指标量",
            float_format="%.5f",
        )
        washed_pd_pearson_cc_matrix.to_excel(
            writer,
            sheet_name="任务2_皮尔逊相关系数",
            header=washed_feature_name_list,
            index_label="指标量",
            float_format="%.5f",
        )

    logger.info("Done")
